module Courcoubetis where

import Math.LinearEquationSolver
import Data.Array	
import Data.List
import Data.Ratio
import System.IO.Unsafe (unsafePerformIO)

----------------------------------------------------
----------------------------------------------------
-- Data type for DTMCs
data DTMC a = DTMC {
			statesDTMC	:: Array Int a,		    -- States
			pTransDTMC	:: Array (Int, Int) Float,  -- State Transition probabilities
			labelFDTMC	:: Array Int [String],	    -- Labeling Function
			initDiDTMC      :: Array Int Float	    -- Initial Distribution Matrix
			} deriving (Show)

-- Values of a sample DTMC
myMatrix :: Array (Int, Int) Float
myMatrix = array ((1,1),(6,6)) 
			 [((1,1), 0.07), ((1,2), 0.28), ((1,3), 0.35), ((1,4), 0.21), ((1,5),0.06), ((1,6),0.03),
			  ((2,1), 0.07), ((2,2), 0.28), ((2,3), 0.35), ((2,4), 0.21), ((2,5),0.06), ((2,6),0.03),
			  ((3,1), 0.07), ((3,2), 0.28), ((3,3), 0.35), ((3,4), 0.21), ((3,5),0.06), ((3,6),0.03),
			  ((4,1), 0.04), ((4,2), 0.16), ((4,3), 0.2), ((4,4), 0.42), ((4,5),0.12), ((4,6),0.06),
			  ((5,1), 0.04), ((5,2), 0.16), ((5,3), 0.2), ((5,4), 0.42), ((5,5),0.12), ((5,6),0.06),
			  ((6,1), 0.04), ((6,2), 0.16), ((6,3), 0.2), ((6,4), 0.42), ((6,5),0.12), ((6,6),0.06)
			  ]

estados :: Array Int (Courcou (Int, Char))
estados = array (1,6) [(1, Base (1,'S')), (2, Base (2,'M')), (3, Base (3,'L')), (4, Base (4,'S')), (5, Base (5,'M')), (6, Base (6,'L'))]

etiquetas :: Array Int [String]
etiquetas = array (1,6) [(1,["p","S"]), (2,["p","M"]), (3,["p","L"]), (4,["q","S"]), (5,["q","M"]), (6,["q","L"])]

inicial :: Array Int Float
inicial = array (1,6) [(1,0.06),(2,0.24),(3,0.3),(4,0.28),(5,0.08),(6,0.04)]

-- An instance of a DTMC named mChain
mChain = DTMC {statesDTMC = estados, pTransDTMC = myMatrix, labelFDTMC = etiquetas, initDiDTMC = inicial}

----------------------------------------------------
----------------------------------------------------
-- Data type for the DTMC states that we obtain by transforming the HMM
data Courcou a = Base a | Par (Courcou a, [Char])

-- Instantiating the Show class for Courcou data type
instance (Show a) => Show (Courcou a) where
	show = showC

showC :: Show a => Courcou a -> String
showC (Base n) = show n
showC (Par (c, s)) = "(" ++ showC c ++ "," ++ show s ++ ")"

-- Instanciating the Eq class for Courcou data type
instance (Eq a) => Eq (Courcou a) where
	(==) = eqC

eqC :: Eq a => (Courcou a) -> (Courcou a) -> Bool
eqC (Base n) (Base m) 		= n==m
eqC (Base _) _ 			= False
eqC _ (Base _)			= False
eqC (Par (c1,s1)) (Par (c2,s2))	= (eqC c1 c2) && (s1 == s2)

-- This function allows us to get the first component of an
-- element of type Courcou a
first :: Courcou a -> Courcou a
first (Base _) 		= error "There is no first for this pattern"
first (Par (f,_)) 	= f

-- This functions allows us to get the second component of an
-- element of type Courcou a
second :: Courcou a -> [Char]
second (Base _)		= error "There is no second for this pattern"
second (Par (_,s))  = s

----------------------------------------------------
----------------------------------------------------
-- Data type for the LTL modal temporal logic
data LTL  = Verdad | Falso | Atom String | No LTL | O LTL LTL | Y LTL LTL
			| Next LTL | UntilB Int LTL LTL | Until LTL LTL

-- Instantiating the Show class for the LTL data type			
instance Show LTL where
	show = ltlStr

ltlStr :: LTL -> String
ltlStr Verdad		= "T"
ltlStr Falso		= "F"
ltlStr (Atom str) 	= str
ltlStr (No p) 		= "~(" ++ ltlStr p ++ ")"
ltlStr (O p1 p2) 	= (ltlStr p1) ++ " v " ++ (ltlStr p2)
ltlStr (Y p1 p2) 	= (ltlStr p1) ++ " ^ " ++ (ltlStr p2)
ltlStr (Next p) 	= "X( " ++ ltlStr p ++ " )"
ltlStr (UntilB n p1 p2) = (ltlStr p1) ++ " U[<=" ++ show n ++ "] " ++ (ltlStr p2)
ltlStr (Until p1 p2)    = (ltlStr p1) ++ " U " ++ (ltlStr p2)

----------------------------------------------------
----------------------------------------------------
-- The Cx and Cu construction. 
-- Here we will get the new M' matrix
cxu :: Eq a => DTMC (Courcou (a, Char)) -- The original DTMC that is going to be extended 
                -> LTL                  -- The LTL formula from which the construction is based
                -> Int                  -- The number for the new variable
                -> (DTMC (Courcou (a, Char)), Int)
cxu m (Next phi) n 	= (DTMC {statesDTMC = newEdos, 
                                 pTransDTMC = newTransition p newM newEdos xY xN xQ arreg, 					
                                 labelFDTMC = newLabeling newM newEdos, 					
                                 initDiDTMC = newInitDist p newM newEdos xY xN arreg}	
                          , nextN)
 				where 					
                                  (newM, nextN)	= cxu m phi (n+1) 					
                                  (p, _)	= reWrite phi (n+1) 					
                                  xY            = xYes p newM
                                  xN            = xNo p newM
                                  xQ            = xQM p newM xY xN
                                  arreg		= arProb newM p xY xN                          			
                                  newEdos 	= newStates p newM n xY xN xQ 

cxu m (UntilB i phi1 phi2) n= (DTMC {statesDTMC = newEdos,					
                                     pTransDTMC = newTransitionU p1 p2 newM newEdos xYu xNu xQu arreg,					
                                     labelFDTMC = newLabelingU newM newEdos,					
                                     initDiDTMC = newInitDistU p1 p2 newM newEdos xYu xNu arreg}
			      , nextN)
				where 	                                	
                                  (tmpM, tmpN)	= cxu m phi1 (n+1)
                                  (newM, nextN)	= cxu tmpM phi2 tmpN								
                                  (p1, nn)	= reWrite phi1 (n+1)					
                                  (p2, _)	= reWrite phi2 nn								
                                  arreg		= arProbU newM p1 p2 i
                                  xYu           = xYesU newM arreg
                                  xNu           = xNoU newM arreg
                                  xQu           = xQMU newM xYu xNu
                                  newEdos	= newStatesU p1 p2 newM n xYu xNu xQu

cxu m (Until phi1 phi2) n= (DTMC {statesDTMC = newEdos,					
                                  pTransDTMC = newTransitionU p1 p2 newM newEdos sY sN sQ arreg,					
                                  labelFDTMC = newLabelingU newM newEdos,					
                                  initDiDTMC = newInitDistU p1 p2 newM newEdos sY sN arreg}
			   , nextN)
				where 	                                	
                                  (tmpM, tmpN)	= cxu m phi1 (n+1)
                                  (newM, nextN)	= cxu tmpM phi2 tmpN								
                                  (p1, nn)	= reWrite phi1 (n+1)					
                                  (p2, _)	= reWrite phi2 nn								
                                  sN            = sNoU p1 p2 newM
                                  sY            = sYesU p1 p2 sN newM
                                  sQ            = sQMU newM sY sN
                                  arreg		= arProbUU newM sY sQ
                                  newEdos	= newStatesU p1 p2 newM n sY sN sQ

cxu m (No p)	n 	= cxu m p n
cxu m (O p1 p2) n 	= cxu newM p2 nextN
				where (newM, nextN) = cxu m p1 n
cxu m (Y p1 p2) n 	= cxu newM p2 nextN
				where (newM, nextN) = cxu m p1 n
cxu m _  n		= (m,n)	

-- And the new formula f' is got with the function reWrite that appears next
reWrite :: LTL -> Int -> (LTL, Int)
reWrite Verdad n	= (Verdad, n)
reWrite Falso n		= (Falso , n)
reWrite (Atom s) n	= (Atom s, n)
reWrite (No p) n	= (No  pp, nextN)
	where (pp, nextN) = reWrite p n
reWrite (O p1 p2) n	= (O pp1 pp2, nextN2)
	where
		(pp1, nextN1) = reWrite p1 n
		(pp2, nextN2) = reWrite p2 nextN1
reWrite (Y p1 p2) n	= (Y pp1 pp2, nextN2)
	where
		(pp1, nextN1) = reWrite p1 n
		(pp2, nextN2) = reWrite p2 nextN1
reWrite (Next p) n      = (Atom ("xi"++show n), nextN)
	where (_, nextN) = reWrite p (n+1)
reWrite (UntilB _ p1 p2) n = (Atom ("xi"++show n), nextN)
	where 
		(_, tmpValue)	= reWrite p1 (n+1)					
                (_, nextN) 	= reWrite p2 tmpValue
reWrite (Until p1 p2) n = (Atom ("xi"++show n), nextN)
	where 
		(_, tmpValue)	= reWrite p1 (n+1)					
                (_, nextN) 	= reWrite p2 tmpValue

----------------------------------------------------
----------------------------------------------------
-- This function gets the observation that is the second component
-- of the state produced by transforming the HMM into a DTMC
getObs :: Eq a => Courcou (a,Char) -> Char
getObs (Base (_, o)) = o
getObs (Par (a, _)) = getObs a

----------------------------------------------------
----------------------------------------------------
-- This functions computes the set of states that satisfy the LTL formula
-- Notice that sat only handles the case when the LTL formula is propositional. 
-- The temporal operators are not considered since this function is invoking
-- on formulas that are temporal operator free. Specifically, we identify the 
-- more nested temporal operator and apply the sat function to its argument
sat :: Eq a => DTMC (Courcou (a, Char)) -> LTL -> [Courcou (a, Char)]
sat m Verdad	= elems (statesDTMC m)
sat _ Falso	= []
sat m (Atom s)	= [ edos!ind | ind<-indices labl, elem s (labl!ind) || elem (getObs (edos!ind)) s]
	where
		labl = labelFDTMC m			
                edos = statesDTMC m
sat m (No p) 	= [ edos!ind | ind<-indices edos, not (elem (edos!ind) (sat m p))]
	where 
		edos = statesDTMC m
sat m (O p1 p2) = union (sat m p1) (sat m p2)
sat m (Y p1 p2) = intersect (sat m p1) (sat m p2)
sat m _		= error "The function sat does not knows how to handle temporal operators"

----------------------------------------------------
----------------------------------------------------
-- The partioning of the set of states in computed next
-- XYes
-- Notice that there must be at leat one transition that has positive
-- probability since the trajectories are infinite
xYes :: Eq a => LTL -> DTMC (Courcou (a, Char)) -> [Courcou (a, Char)]
xYes phi mC = [edos!i | i<-indices edos, checkAll i]
	where
		edos = statesDTMC mC			
                checkAll ind = and [ if (pTransDTMC mC)!(ind, j) > 0 
					then elem (edos!j) (sat mC phi) 
					else True
					| j<-indices edos]

-- XNo
xNo :: Eq a => LTL -> DTMC (Courcou (a, Char)) -> [Courcou (a, Char)]
xNo phi mC = [edos!i | i<-indices edos, checkAll i]
	where
		edos = statesDTMC mC			
                checkAll ind = and [ if (pTransDTMC mC)!(ind, j) > 0 
					then not(elem (edos!j) (sat mC phi))
					else True
					| j<-indices edos]

-- X?
xQM :: Eq a => LTL 
               -> DTMC (Courcou (a, Char)) 
               -> [Courcou (a, Char)]
               -> [Courcou (a, Char)]
               -> [Courcou (a, Char)]
xQM phi mC xY xN = elems (statesDTMC mC) \\ (xY ++ xN)

----------------------------------------------------
----------------------------------------------------
-- We compute the new states with this function
newStates :: Eq a => LTL                                -- The temporal formula that is being replaced
                     -> DTMC (Courcou (a, Char))        -- The current DTMC 
                     -> Int                             -- The number for the new variable that will replaced the LTL formula
                     -> [Courcou (a, Char)]             -- The set xYes
                     -> [Courcou (a, Char)]             -- The set xNo
                     -> [Courcou (a, Char)]             -- The set xQ
                     -> Array Int (Courcou (a, Char))   -- The resulting array with the new states
newStates phi mC n yes no qMark = array (1, length asoccL) (zip [1..] asoccL)
	where
		edos   = elems (statesDTMC mC)				
                asoccL = [ if elem s yes 
				then Par (s,"xi"++show n) 
				else if elem s no 
					then Par (s, "noXi"++show n) 
					else Par (s, "xi"++show n) | s<-edos]
				++ [Par (s, "noXi"++show n) | s<-edos, elem s qMark]							

----------------------------------------------------
----------------------------------------------------
-- Function that computes the new initial distribution
newInitDist :: Eq a => LTL                                -- The temporal formula that is being replaced
                       -> DTMC (Courcou (a, Char))        -- The current DTMC
                       -> Array Int (Courcou (a, Char))   -- The new set of states
                       -> [Courcou (a, Char)]             -- The set xYes
                       -> [Courcou (a, Char)]             -- The set xNo
                       -> Array Int Float                 -- The array with the probabilities for each state to satisfy the temporal formula
                       -> Array Int Float                 -- The resulting array with the new initial distribution
newInitDist phi mC newStat xY xN arreg = array (bounds newStat) [(i,computeInitDis (getInd mC (first(newStat!i))) (newStat!i)) 
								| i<-(indices newStat)]
		where
			computeInitDis indU (Par (u,x)) | (elem u xY) || (elem u xN)	=  initDist!indU
							| (take 2 x /= "no")  		= (initDist!indU)*(arreg!indU)
							| otherwise			= (initDist!indU)*(1-(arreg!indU))
			computeInitDis _ _ = error "This case shouldn't happen"					
                        initDist= initDiDTMC mC	

----------------------------------------------------
----------------------------------------------------
-- Now we update the labelling function
newLabeling :: Eq a => DTMC (Courcou (a, Char))         -- The current DTMC 
                       -> Array Int (Courcou (a, Char)) -- The new set of states
                       -> Array Int [String]            -- The resulting array with the new labelling function for each state
newLabeling mC newStat = array (bounds newStat) [(i, getAtom i) | i<-indices newStat]
		where
			getAtom ind = [str | str<-(labelFDTMC mC)!(getInd mC (first (newStat!ind)))] 
					++ [second (newStat!ind)]

----------------------------------------------------
----------------------------------------------------
-- We compute the new state transition matrix
newTransition :: Eq a => LTL                              -- The temporal formula that is being replaced
                         -> DTMC (Courcou (a, Char))      -- The current DTMC 
                         -> Array Int (Courcou (a, Char)) -- The new set of states 
                         -> [Courcou (a, Char)]           -- The set xYes
                         -> [Courcou (a, Char)]           -- The set xNo
                         -> [Courcou (a, Char)]           -- The set xQ
                         -> Array Int Float               -- The array with the probabilities for each state to satisfy the temporal formula
                         -> Array (Int, Int) Float        -- The resulting array with the new transition probabilities
newTransition phi mC newStat xY xN xQ arreg = array ((1,1), (sb, sb)) lista
		where
			sb = snd (bounds newStat)				
                        lista = [ ((i,j), getVal mC phi (newStat!i) (newStat!j) xY xN xQ arreg)
				| i<-[1..sb], j<-[1..sb]]

----------------------------------------------------
----------------------------------------------------
-- Auxiliary functions
-- This one computes the value for each matrix entry
getVal :: Eq a => DTMC (Courcou (a, Char))           -- The current DTMC 
                  -> LTL                             -- The temporal formula that is being replaced
                  -> Courcou (a, Char)               -- The start state from which the transition is defined 
                  -> Courcou (a, Char)               -- The destination state to which the transition is defined
                  -> [Courcou (a, Char)]             -- The set xYes
                  -> [Courcou (a, Char)]             -- The set xNo
                  -> [Courcou (a, Char)]             -- The set xQ
                  -> Array Int Float                 -- The array with the probabilities for each state to satisfy the temporal formula
                  -> Float                           -- The value for this transition
getVal mC phi (Par (u,x)) (Par (v,y)) xY xN xQ ar
  | (elem u xY || elem u xN)= if (elem v xY || elem v xN) 
                              then pTran!(gu, gv)
                              else if (take 2 y /= "no") 
                                   then (pTran!(gu, gv))*(gQv)
                                   else (pTran!(gu, gv))*(1-gQv) 
  | elem v xQ		    = if (take 2 x /= "no")
                              then if vSatPhi
                                   then if (take 2 y /= "no") 
                                        then (pTran!(gu, gv))*(gQv)/(gQu) 
                                        else (pTran!(gu, gv))*(1-gQv)/(gQu)
                                   else 0
                              else if (not vSatPhi)
                                   then if (take 2 y /= "no") 
                                        then (pTran!(gu, gv))*(gQv)/(1-gQu) 
                                        else (pTran!(gu, gv))*(1-gQv)/(1-gQu)
                                   else 0
  | (elem v xY || elem v xN)= if vSatPhi
                              then if (take 2 x /= "no") 
                                   then (pTran!(gu, gv))/(gQu) 
                                   else 0
                              else if (take 2 x == "no") 
                                   then (pTran!(gu, gv))/(1-gQu) 
                                   else 0
  | otherwise		    = 0
  where 
    gu = getInd mC u                        
    gv = getInd mC v                             
    gQu = ar!gu                                   
    gQv = ar!gv                                         
    pTran = pTransDTMC mC                                                 
    vSatPhi = elem v (sat mC phi)

-- This method helps us to determine the state's index in a Marcov chain
getInd :: Eq a => DTMC (Courcou (a, Char)) -> Courcou (a, Char) -> Int
getInd mC s = case elemIndex s sts of 
                Just x    -> x+1
                otherwise -> error "Error trying to get index of element"
	where sts = elems (statesDTMC mC)

-- Given a matrix this function returns the ith row
row :: (Ix a, Ix b) => a -> Array (a,b) c -> Array b c
row i m = ixmap (l',u') (\j->(i,j)) m
	where ((_,l'),(_,u')) = bounds m

-- This function computes the values (qu) needed by the construction cxu and stores them in an array
arProb :: Eq a => DTMC(Courcou (a, Char))            -- The current DTMC
                  -> LTL                             -- The LTL formula that is operand of the next operator
                  -> [Courcou (a, Char)]             -- The set xYes
                  -> [Courcou (a, Char)]             -- The set xNo
                  -> Array Int Float
arProb mC phi yes no = listArray (bounds st) [if elem (st!i) yes
                                              then 1
                                              else if elem (st!i) no
                                                   then 0
                                                   else sum [ (row i (pTransDTMC mC))!j | 
                                                              j<-(map (getInd mC) (filter (\x->elem x (sat mC phi)) (elems st)))]
                                             | i<-indices st]
        where       
          st = statesDTMC mC



--------------------------------------------------------------------------------------------------------
--------------------------------------------------------------------------------------------------------
-- The Cu construction appears next. The bounded version.
--------------------------------------------------------------------------------------------------------
--------------------------------------------------------------------------------------------------------

-- The partioning of the set of states for the construction of Cu in computed next
-- XYes
xYesU :: Eq a => DTMC (Courcou (a, Char)) -> Array Int Float -> [Courcou (a, Char)]
xYesU mC ar = [(st!i) | i<-indices st, (ar!i)==1] 
  where 
    st = statesDTMC mC

-- XNo
xNoU :: Eq a => DTMC (Courcou (a, Char)) -> Array Int Float -> [Courcou (a, Char)]
xNoU mC ar = [(st!i) | i<-indices st, (ar!i)==0] 
  where 
    st = statesDTMC mC

-- X?										
xQMU :: Eq a => DTMC (Courcou (a, Char)) 
                -> [Courcou (a, Char)]
                -> [Courcou (a, Char)]
                -> [Courcou (a, Char)]
xQMU mC xYU xNU = elems (statesDTMC mC) \\ (xYU ++ xNU)

----------------------------------------------------
----------------------------------------------------
-- We compute the new states with this function
newStatesU :: Eq a => LTL                               -- The LHS formula operand of the until operator 
                      -> LTL                            -- The RHS formula operand of the until operator
                      -> DTMC (Courcou (a, Char))       -- The current DTMC 
                      -> Int                            -- The number for the new variable that will replaced the LTL formula
                      -> [Courcou (a, Char)]            -- The set xYes
                      -> [Courcou (a, Char)]            -- The set xNo
                      -> [Courcou (a, Char)]            -- The set xQ
                      -> Array Int (Courcou (a, Char))  -- The resulting array with the new states
newStatesU phi1 phi2 mC n yes no qMark = array (1, length asoccL) (zip [1..] asoccL)
        where
                edos   = elems (statesDTMC mC)
                asoccL = [ if elem s yes                         
                           then Par (s,"xi"++show n) 
                           else if elem s no                     
                                then Par (s, "noXi"++show n)                     
                                else Par (s, "xi"++show n) | s<-edos]             
                         ++ [Par (s, "noXi"++show n) | s<-edos,                  
                             elem s qMark]							

----------------------------------------------------
----------------------------------------------------
-- Function that computes the new initial distribution
-- Check this one!!! Probably the n is not used
newInitDistU :: Eq a => LTL                               -- The LHS formula operand of the until operator 
                        -> LTL                            -- The RHS formula operand of the until operator
                        -> DTMC (Courcou (a, Char))       -- The current DTMC
                        -> Array Int (Courcou (a, Char))  -- The new set of states
                        -> [Courcou (a, Char)]            -- The set xYes
                        -> [Courcou (a, Char)]            -- The set xNo
                        -> Array Int Float                -- The array with the probabilities for each state to satisfy the temporal formula
                        -> Array Int Float                -- The resulting array with the new initial distribution
newInitDistU phi1 phi2 mC newStat xY xN arreg = array (bounds newStat) 
                                          [(i,computeInitDis (getInd mC (first(newStat!i))) (newStat!i)) | i<-(indices newStat)]
        where                                  
          computeInitDis indU (Par (u,x)) | (elem u xY) || (elem u xN)	=  initDist!indU
                                          | (take 2 x /= "no")		= (initDist!indU)*(arreg!indU)
                                          | otherwise			= (initDist!indU)*(1-(arreg!indU))
          initDist = initDiDTMC mC 
----------------------------------------------------
----------------------------------------------------
-- Now we update the labeling function
newLabelingU :: Eq a => DTMC (Courcou (a, Char))        -- The current DTMC 
                        -> Array Int (Courcou (a, Char))-- The new set of states 
                        -> Array Int [String]           -- The resulting array with the new labelling function for each state
newLabelingU mC newStat = array (bounds newStat) [(i, getAtom i) | i<-indices newStat]
          where								
            getAtom ind = [str | str<-(labelFDTMC mC)!(getInd mC (first (newStat!ind)))] 				
                          ++ [second (newStat!ind)]


----------------------------------------------------
----------------------------------------------------
-- We compute the new state transition matrix for Cu
newTransitionU :: Eq a => LTL                                   -- The LHS formula operand of the until operator
                          -> LTL                                -- The RHS formula operand of the until operator
                          -> DTMC(Courcou (a, Char))            -- The current DTMC
                          -> Array Int (Courcou (a, Char))      -- The new set of states
                          -> [Courcou (a, Char)]                -- The set xYes
                          -> [Courcou (a, Char)]                -- The set xNo
                          -> [Courcou (a, Char)]                -- The set xQ
                          -> Array Int Float                    -- The array with the probabilities for each state to satisfy the temporal formula
                          -> Array (Int, Int) Float             -- The resulting array with the new transition probabilities
newTransitionU phi1 phi2 mC newStat xY xN xQ arreg = array ((1,1), (sb, sb)) lista
          where	
            sb     = snd (bounds newStat)		
            lista  = [ ((i,j), getValU mC phi1 phi2 (newStat!i) (newStat!j) xY xN xQ arreg)
			| i<-[1..sb], j<-[1..sb]]

----------------------------------------------------
----------------------------------------------------
-- Auxiliary functions
-- This one computes the value for each transition probability matrix entry
getValU :: Eq a => DTMC (Courcou (a, Char))           -- The current DTMC
                   -> LTL                             -- The LHS formula operand of the until operator
                   -> LTL                             -- The RHS formula operand of the until operator
                   -> Courcou (a, Char)               -- The start state from which the transition is defined 
                   -> Courcou (a, Char)               -- The destination state to which the transition is defined 
                   -> [Courcou (a, Char)]             -- The set xYes
                   -> [Courcou (a, Char)]             -- The set xNo
                   -> [Courcou (a, Char)]             -- The set xQ
                   -> Array Int Float
                   -> Float
getValU mC phi1 phi2 (Par (u,x)) (Par (v,y)) xY xN xQ ar 
  | (elem u xY || elem u xN)	= if (elem v xY || elem v xN) 
                                  then pTran!(gu, gv)										
                                  else if (take 2 y /= "no")						
                                       then (pTran!(gu, gv))*(gQv)						
                                       else (pTran!(gu, gv))*(1-gQv) 					  			
  | elem v xQ 			= if (take 2 x /= "no") 										
                                  then if (take 2 y /= "no") 										
                                       then (pTran!(gu, gv))*(gQv)/(gQu) 							
                                       else 0	 		 						
                                  else if (take 2 y == "no")									
                                       then (pTran!(gu, gv))*(1-gQv)/(1-gQu)							
                                       else 0									 	
  | elem v xY			= if (take 2 x /= "no") 										
                                  then (pTran!(gu, gv))/(gQu)									
                                  else 0							  			
  | elem v xN			= if (take 2 x == "no") 							
                                  then (pTran!(gu, gv))/(1-gQu) 						
                                  else 0				
  | otherwise			= 0	
  where 
    gu = getInd mC u
    gv = getInd mC v
    gQu = ar!gu
    gQv = ar!gv
    pTran = pTransDTMC mC

-- This function computes the values (qu) needed by the construction cu and stores them in an array
arProbU :: Eq a => DTMC(Courcou (a, Char))            -- The current DTMC
                   -> LTL                             -- The LHS formula operand of the until operator
                   -> LTL                             -- The RHS formula operand of the until operator
                   -> Int                             -- The integer that bounds the until operator
                   -> Array Int Float                 -- The array with the probabilities that each state satisfies the until formula
arProbU mC phi1 phi2 m = listArray (1, length myList) myList                    
  where                           
    myList = auxArProbU mC m yes qmark
    elemen = elems (statesDTMC mC)
    yes    = sat mC phi2
    no     = elemen \\ (union (sat mC phi1) yes) 
    qmark  = elemen \\ (yes ++ no)

-- Auxiliary function that computes the probability that (phi1 U phi2) is satisfies
-- in each state
auxArProbU :: Eq a => DTMC(Courcou (a, Char))            -- The current DTMC
                      -> Int                             -- The integer that bounds the until operator
                      -> [Courcou (a, Char)]             -- The set xYes
                      -> [Courcou (a, Char)]             -- The set xQ
                      -> [Float]                         -- The list with the probabilities that each state satisfies the until formula
auxArProbU mC 0 sY _  = [if elem edo sY then 1.0 else 0.0 | edo<-elems (statesDTMC mC)]
auxArProbU mC m sY sQ = mult pPrime (auxArProbU mC (m-1) sY sQ)			
  where 
    mult xss ys = [sum (zipWith (*) xs ys) | xs<-xss] 
    pPrime      = [takeVal s (elems (row (getInd mC s) (pTransDTMC mC))) 
                  | s<-elems (statesM)]
    statesM     = statesDTMC mC								 
    takeVal st list  = if (elem st sQ) 
                        then list
                        else if (elem st sY)
                             then [if e==st then 1.0 else 0.0 | e<-elems statesM]
                             else replicate (length list) 0.0



--------------------------------------------------------------------------------------------------------
--------------------------------------------------------------------------------------------------------
-- The Cu construction appears next. (The unbounded version)
--------------------------------------------------------------------------------------------------------
--------------------------------------------------------------------------------------------------------
-- The partioning of the set of states for the construction of Cu in computed next
-- Set No
-- The set for which the p_s(p U q) is exactly 0
sNoU :: Eq a => LTL -> LTL -> DTMC (Courcou (a, Char)) -> [Courcou (a, Char)]
sNoU phi1 phi2 mC = (elems (statesDTMC mC)) \\ (prob0 mC (sat mC phi1) (sat mC phi2))

-- An auxiliary function to compute the fixed point need to obtain the set No 
prob0 :: Eq a => DTMC (Courcou (a, Char)) 
                 -> [Courcou (a, Char)] 
                 -> [Courcou (a, Char)] 
                 -> [Courcou (a, Char)] 
prob0 mC sat1 sat2 = if sat2 == nL
                         then sat2 
                         else prob0 mC sat1 nL
  where
    nL = union sat2 [st!i | i<-indices st,             
                     elem (st!i) sat1,
                     or (map (\(i,x) -> x>0 && elem (st!i) sat2) (assocs (row i (pTransDTMC mC))))]
    st = statesDTMC mC
    
-- Set Yes
-- The set for which the p_s(p U q) is exactly 1
sYesU :: Eq a => LTL -> LTL -> [Courcou (a, Char)] -> DTMC (Courcou (a, Char)) -> [Courcou (a, Char)]
sYesU phi1 phi2 sN mC = (elems (statesDTMC mC)) \\ (prob1 mC ((sat mC phi1) \\ (sat mC phi2)) sN)

-- An auxiliary function to compute the fixed point need to obtain the set Yes 
prob1 :: Eq a =>  DTMC (Courcou (a, Char)) 
                 -> [Courcou (a, Char)] 
                 -> [Courcou (a, Char)] 
                 -> [Courcou (a, Char)] 
prob1 mC diff sN  = if sN == nL                       
                       then sN                         
                       else prob1 mC diff nL
  where
    nL = union sN [st!i | i<-indices st,             
                     elem (st!i) diff,
                     or (map (\(i,x) -> x>0 && elem (st!i) sN) (assocs (row i (pTransDTMC mC))))]
    st = statesDTMC mC

-- Set ?
-- The states that are neither in sYesU nor in sNoU
sQMU :: Eq a => DTMC (Courcou (a, Char)) 
         -> [Courcou (a, Char)] 
         -> [Courcou (a, Char)] 
         -> [Courcou (a, Char)] 
sQMU mC sY sN = (elems (statesDTMC mC)) \\ (sY ++ sN)

-- In order to find the probability that a state in set ? satisfies the 
-- formula p_s (p U q) we need to solve a linear equation system (LES) of the  
-- form Ax=b, where A=I-(a modified version of the transition probability matrix) 
-- and b is a column with 1 in the i-row if the i state is in the set Yes, and 0 otherwise.
-- The next function computes the matrix A involved in the LES.
aMat :: Eq a => DTMC (Courcou (a, Char)) 
                  -> [Courcou (a, Char)] 
                  -> [[Rational]]
aMat mC sQ = [if elem (st!i) sQ 
                then (rowPrime i trans)
                else [if j==i 
                      then 1.0 
                      else 0.0| j<-ind] 
                | i<-ind]
  where
    st    = statesDTMC mC
    trans = pTransDTMC mC
    ind   = indices st
    
-- This is an auxiliary function that returns the i column of the
-- previous A matrix. Notice how in the diagonal we substract 
-- from 1 the probability P(s,s), elsewhere we have the probability 
-- negated, i.e. -P(s,s')
rowPrime :: (Ix a, Real b) => a -> Array (a,a) b -> [Rational]
rowPrime i m = [if k==i 
                then toRational (1 - (m!(i,k)))             
                else toRational (  - (m!(i,k))) | k<-range (l',u')] 
  where ((_,l'),(_,u')) = bounds m
        
-- b is a column with 1 in the row with corresponding state in the set Yes, and 0 otherwise.
bCol :: Eq a => DTMC (Courcou (a, Char)) 
                  -> [Courcou (a, Char)]           
                  -> [Rational]
bCol mC sY = [if elem st sY 
                then 1 
                else 0 | st<-elems (statesDTMC mC)]
             
--  This function computes the probabilities that each state satisfies the until formula           
arProbUU :: Eq a =>  DTMC (Courcou (a, Char))    -- The current DTMC
                     -> [Courcou (a, Char)]      -- The set Yes
                     -> [Courcou (a, Char)]      -- The set ?
                     -> Array Int Float          -- The array with the probabilities that each state satisfies the until formula
arProbUU mC sY sQ = listArray (1, length myL) myL 
  where
    myL = case unsafePerformIO (solveRationalLinearEqs z3 (aMat mC sQ) (bCol mC sY)) of 
                        Just xs   -> map (\r->fromRational r) xs 
                        otherwise -> error "There was an error trying to get the solution of the linear equation system"
